'''
清洗完数据以后，我们要开始训练模型并对训练出的模型精度进行评估
这里是先验为多项式分布的朴素贝叶斯模型
'''

from pymongo import MongoClient
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
from sklearn.naive_bayes import MultinomialNB
from sklearn import metrics
import sys


if __name__ == "__main__":
    # 先连接数据库
    client = MongoClient()
    db = client["LawData"]

    # 读入训练数据和测试数据
    print("loading data...")
    train_text = []
    train_label = []
    test_text = []
    test_label = []
    for item in db.train_jieba.find():
        train_text.append(item["content"])
        train_label.append(item["label"])
    for item in db.test_jieba.find():
        test_text.append(item["content"])
        test_label.append(item["label"])
    all_text = train_text + test_text
    print("loading finish")

    # 将文本转换为向量
    print("text to vec...")
    count_v0 = CountVectorizer()
    # 计算所有词的出现次数
    counts_all = count_v0.fit_transform(all_text)
    # 将词语转换为词频矩阵
    count_v1 = CountVectorizer(vocabulary=count_v0.vocabulary_)
    counts_train = count_v1.fit_transform(train_text)
    # print("The shape of train is ", repr(counts_train.shape))
    count_v2 = CountVectorizer(vocabulary=count_v0.vocabulary_)
    counts_test = count_v2.fit_transform(test_text)
    # print("The shape of test is ", repr(counts_test.shape))
    # 将词频矩阵转换为tf-idf矩阵
    tf = TfidfTransformer()
    train_data = tf.fit(counts_train).transform(counts_train)
    test_data = tf.fit(counts_test).transform(counts_test)

    x_train = train_data
    y_train = train_label
    x_test = test_data
    y_test = test_label
    print("text to vec finish")

    # 选择模型为：先验为多项式分布的朴素贝叶斯
    clf = MultinomialNB(alpha=0.01)
    clf.fit(x_train, y_train)
    
    # 以上训练完成后根据我们的测试集数据给出预测值
    predicts = clf.predict(x_test)
    num = 0
    preds = predicts.tolist()
    for i, pred in enumerate(preds):
        if int(pred) == int(y_test[i]):
            num += 1
    print("预测准确度：", float(num) / len(preds))

